package xyj.sparrow.discovery.zk.bean;

import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.retry.RetryForever;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceDiscoveryBuilder;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.ServiceInstanceBuilder;
import org.apache.curator.x.discovery.ServiceCache;
import org.springframework.stereotype.Component;
import xyj.sparrow.common.discovery.SparrowServiceInstance;
import xyj.sparrow.discovery.api.SparrowDiscovery;
import xyj.sparrow.discovery.zk.bean.strategies.*;
import xyj.sparrow.discovery.zk.config.ZkConf;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * @author 江南小俊
 * @since 2021/6/19 上午9:56
 **/
@Component
@Slf4j
public class ZkSparrowDiscovery implements SparrowDiscovery {
    @Resource
    private ZkConf zkConf;
    private CuratorFramework client;
    private ServiceDiscovery<Object> discovery;
    private final ConcurrentHashMap<String, ServiceListener> serviceCacheMap =
            new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, ServiceInstance<Object>> serviceInstanceMap = new ConcurrentHashMap<>();

    @PostConstruct
    private void init() throws Exception {
        log.info("ZK init.....");
        //永久重连 间隔1s
        RetryForever retryForever = new RetryForever(1000);
        String zookeeperConnectionString = zkConf.getAddress() + ":" + zkConf.getPort();
        client = CuratorFrameworkFactory.newClient(zookeeperConnectionString, retryForever);
        client.start();
        log.info("ZK  client start.....");
        discovery = ServiceDiscoveryBuilder.builder(Object.class).client(client).basePath(
                "/sparrow").build();
        discovery.start();
        log.info("ZK  discovery start.....");
    }

    @PreDestroy
    private void destroy() {
        log.info("ZK close....");
        client.close();
    }

    /**
     * 分布式锁
     *
     * @param lockPath 锁资源路径
     * @param maxWait  超时解锁
     * @param waitUnit 超时单位
     * @author 江南小俊
     * @since 2021/6/19 上午10:02
     **/

    public void lock(String lockPath, long maxWait, TimeUnit waitUnit) throws Exception {
        InterProcessMutex lock = new InterProcessMutex(client, lockPath);
        if (lock.acquire(maxWait, waitUnit)) {
            try {
                // do some work inside of the critical section here
                log.info("获取分布式锁...");
                Thread.sleep(5000);
            } finally {
                lock.release();
                log.info("释放分布式锁...");

            }
        }
    }

    /**
     * 注册/重新注册服务
     *
     * @param serviceName 服务名
     * @param address     ip 可为空，默认读取网卡ip
     * @param port        端口号
     * @param payload     自定义数据
     * @author 江南小俊
     * @since 2021/6/19 下午10:30
     **/
    @Override
    public Object register(String serviceName, String address, int port, Object payload) throws Exception {
        log.info("service discovery register:{}-{}:{}", serviceName, address, port);
        ServiceInstanceBuilder<Object> builder = ServiceInstance.builder();
        if (!StrUtil.isBlank(address)) {
            builder.address(address);
        }
        ServiceInstance<Object> service =
                builder.port(port).payload(payload).name(serviceName).build();
        discovery.registerService(service);
        serviceInstanceMap.put(service.getId(), service);
        return service.getId();
    }

    /**
     * 取消注册/删除服务实例
     *
     * @param serviceIds 服务ids
     * @author 江南小俊
     * @since 2021/6/19 下午10:30
     **/
    @Override
    public void unregister(List<String> serviceIds) {
        log.info("service discovery unregister:{}", serviceIds);
        serviceIds.forEach(serviceId -> {
            try {
                discovery.unregisterService(serviceInstanceMap.get(serviceId));
            } catch (Exception e) {
                e.printStackTrace();
                log.error("", e);
            }
        });
    }


    @Override
    public void addListener(String serviceName, String providerStrategy, Consumer consumer) throws Exception {
        ServiceCache<Object> serviceCache = discovery.serviceCacheBuilder().name(serviceName).build();
        List<ServiceInstance<Object>> instances = serviceCache.getInstances();
        ServiceListener serviceCacheListener;
        switch (providerStrategy) {
            case ProviderStrategy.RANDOM:
                serviceCacheListener = new RandomServiceListener(serviceCache, instances, serviceName, consumer);

                break;
            case ProviderStrategy.STICKY:
                serviceCacheListener = new StickyServiceListener(serviceCache, instances, serviceName, consumer);
                break;
            default:
                serviceCacheListener = new RoundRobinServiceListener(serviceCache, instances, serviceName, consumer);
        }
        serviceCache.addListener(serviceCacheListener);
        serviceCacheMap.put(serviceName, serviceCacheListener);
        serviceCache.start();
        log.info("service：{} addListener.....", serviceName);
    }

    /**
     * 添加服务监听
     *
     * @param serviceName      服务名
     * @param providerStrategy 策略名称  默认{@link ProviderStrategy#ROUND_ROBIN}
     * @author 江南小俊
     * @since 2021/6/19 下午11:02
     **/
    @Override
    public void addListener(String serviceName, String providerStrategy) throws Exception {
        addListener(serviceName, providerStrategy, null);
    }

    @Override
    public void addListener(String serviceName) throws Exception {
        addListener(serviceName, ProviderStrategy.ROUND_ROBIN);
    }

    /**
     * 获取服务实例
     *
     * @param serviceName 服务名
     * @return org.apache.curator.x.discovery.ServiceInstance<java.lang.String> 找不到返回null
     * @author 江南小俊
     * @since 2021/6/19 下午11:21
     **/
    @Override
    public SparrowServiceInstance getInstance(String serviceName) throws Exception {
        if (!serviceCacheMap.containsKey(serviceName)) {
            log.info("缓存没有实例");
            return null;
        }
        ServiceInstance<Object> serviceInstance = serviceCacheMap.get(serviceName).getInstance();
        if (serviceInstance == null) {
            log.info("缓存没有实例");
            return null;
        } else {
            SparrowServiceInstance sparrowServiceInstance = new SparrowServiceInstance();
            sparrowServiceInstance.setName(serviceInstance.getName());
            sparrowServiceInstance.setAddress(serviceInstance.getAddress());
            sparrowServiceInstance.setPort(serviceInstance.getPort());
            sparrowServiceInstance.setPayload(serviceInstance.getPayload());
            return sparrowServiceInstance;
        }
    }


    @Override
    public List<SparrowServiceInstance> listInstance(String serviceName) throws Exception {
        if (!serviceCacheMap.containsKey(serviceName)) {
            log.info("缓存没有实例");
            return new ArrayList<>();
        }
        List<ServiceInstance<Object>> serviceInstances = serviceCacheMap.get(serviceName).listInstance();
        if (serviceInstances == null || serviceInstances.isEmpty()) {
            log.info("缓存没有实例");
            return new ArrayList<>();
        } else {
            List<SparrowServiceInstance> list = new ArrayList<>();
            for (ServiceInstance<Object> serviceInstance : serviceInstances) {
                SparrowServiceInstance sparrowServiceInstance = new SparrowServiceInstance();
                sparrowServiceInstance.setName(serviceInstance.getName());
                sparrowServiceInstance.setAddress(serviceInstance.getAddress());
                sparrowServiceInstance.setPort(serviceInstance.getPort());
                sparrowServiceInstance.setPayload(serviceInstance.getPayload());
                list.add(sparrowServiceInstance);
            }

            return list;
        }
    }

}
